

import math
import calendar
import datetime

from docx.enum.table import (
    WD_CELL_VERTICAL_ALIGNMENT,
    WD_ROW_HEIGHT_RULE
)
from docx.enum.text import (
    WD_PARAGRAPH_ALIGNMENT,
)

from docx.oxml import parse_xml
from docx.oxml.ns import nsdecls, qn
from docx.shared import Mm

from app.libs.display import LEVEL_NUM_STR


class CalendarTable(object):
    """绘制日历威胁表格"""

    def __init__(self, tpl):
        self.tpl = tpl

    @staticmethod
    def _get_calendar_value(item):
        data = {"count": 0, "level": 0}
        for level, count in item.items():
            if level > data["level"] and count:
                data.update({"count": count, "level": level})
        return data

    @staticmethod
    def add_cell_value(cell, context):
        cell.paragraphs[0].alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
        cell.paragraphs[0].add_run(str(context))

        cell.paragraphs[0].paragraph_format.first_line_indent = 0  # 无效， 但是会生成paragraph_format.element.pPr.ind对象
        cell.paragraphs[0].paragraph_format.element.pPr.ind.set((qn("w:firstLineChars")), "0")
        cell.paragraphs[0].paragraph_format.element.pPr.ind.set((qn("w:leftChars")), "0")
        cell.paragraphs[0].paragraph_format.element.pPr.ind.set((qn("w:rightChars")), "0")

        cell.paragraphs[0].paragraph_format.line_spacing = 1
        cell.paragraphs[0].paragraph_format.element.pPr.spacing.set((qn("w:beforeLines")), "0")
        cell.paragraphs[0].paragraph_format.element.pPr.spacing.set((qn("w:before")), "0")
        cell.paragraphs[0].paragraph_format.element.pPr.spacing.set((qn("w:afterLines")), "0")
        cell.paragraphs[0].paragraph_format.element.pPr.spacing.set((qn("w:after")), "0")

        cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER

    @staticmethod
    def get_cell_background_color(level):
        if level == 5:
            return parse_xml(r'<w:shd {} w:fill="C00000"/>'.format(nsdecls('w')))
        elif level == 4:
            return parse_xml(r'<w:shd {} w:fill="FF0000"/>'.format(nsdecls('w')))
        elif level == 3:
            return parse_xml(r'<w:shd {} w:fill="E36C09"/>'.format(nsdecls('w')))
        elif level == 2:
            return parse_xml(r'<w:shd {} w:fill="FFC000"/>'.format(nsdecls('w')))
        return None

    def year_calendar(self, month_list, data):
        tmp_result = {}
        row_count = 9
        col_count = 0
        years = list(sorted({m[0] for m in month_list}))
        for y in years:
            for m in range(1, 13):
                if m < 10:
                    tmp_result.update({f"{y}-0{m}": {level: 0 for level in LEVEL_NUM_STR if level > 1}})
                else:
                    tmp_result.update({f"{y}-{m}": {level: 0 for level in LEVEL_NUM_STR if level > 1}})
            col_count += 3
        for date_str, item in data.items():
            for level, count in item.items():
                tmp_result[date_str[:7]][level] += count
        risk_result = {}
        for date_str, item in tmp_result.items():
            risk_result.update({date_str: self._get_calendar_value(item)})
        sub_doc = self.tpl.new_subdoc()
        table = sub_doc.add_table(rows=row_count, cols=col_count)
        # 年
        for i, y in enumerate(years):
            cell = table.cell(0, i * 3).merge(table.cell(0, (i + 1) * 3 - 1))
            self.add_cell_value(cell, f"{y}年")
        for index, y in enumerate(years):
            for m_index, m in enumerate(range(0, 12)):
                date_str = f"{y}-{m + 1}" if m + 1 >= 10 else f"{y}-0{m + 1}"
                risk_count = risk_result.get(date_str, {}).get("count")
                risk_level = risk_result.get(date_str, {}).get("level")
                col = (index * 3) + m_index % 3
                date_row = 1 + (m_index // 3) * 2
                date_cell = table.cell(date_row, col)
                self.add_cell_value(date_cell, f"{m+1}月")
                shading_elm = parse_xml(r'<w:shd {} w:fill="D9D9D9"/>'.format(nsdecls('w')))
                date_cell._tc.get_or_add_tcPr().append(shading_elm)
                if not risk_count:
                    continue
                risk_row = 2 + (m_index // 3) * 2
                risk_cell = table.cell(risk_row, col)
                self.add_cell_value(risk_cell, str(risk_count))
                color = self.get_cell_background_color(risk_level)
                if color is not None:
                    risk_cell._tc.get_or_add_tcPr().append(color)

        for row in table.rows:
            row.height_rule = WD_ROW_HEIGHT_RULE.EXACTLY
            row.height = Mm(5)
        return sub_doc

    def month_calendar(self, month_list, data):
        risk_data = {}
        for date_str, item in data.items():
            risk_data.update({date_str: self._get_calendar_value(item)})
        month_days = []
        week_str = ["日", "一", "二", "三", "四", "五", "六"]
        rows = 0
        col_count = len(month_list) * 7
        for month_tuple in month_list:
            year, month = month_tuple
            first_day, day_count = calendar.monthrange(year, month)
            if first_day < 6:
                first_day += 1
            else:
                first_day = 0
            days = [0] * first_day + [i for i in range(1, day_count + 1)]
            if math.ceil(len(days) / 7) > rows:
                rows = math.ceil(len(days) / 7)
            month_days.append((month_tuple, days))
        row_count = rows * 2 + 2
        sub_doc = self.tpl.new_subdoc()
        table = sub_doc.add_table(rows=row_count, cols=col_count)
        # 月份
        for i, m in enumerate(month_list):
            cell = table.cell(0, i * 7).merge(table.cell(0, (i + 1) * 7 - 1))
            self.add_cell_value(cell, f"{m[0]}年{m[1]}月")
            shading_elm = parse_xml(r'<w:shd {} w:fill="00FF00"/>'.format(nsdecls('w')))
            cell._tc.get_or_add_tcPr().append(shading_elm)
        # 星期
        for c in range(col_count):
            cell = table.cell(1, c)
            self.add_cell_value(cell, week_str[c % 7])
        for index, month_d in enumerate(month_days):
            for d_index, d in enumerate(month_d[1]):
                if not d:
                    continue
                date_str = datetime.datetime(year=month_d[0][0], month=month_d[0][1], day=d).strftime("%Y-%m-%d")
                risk_count = risk_data.get(date_str, {}).get("count", 0)
                risk_level = risk_data.get(date_str, {}).get("level")

                col = (index * 7) + d_index % 7
                date_row = 2 + (d_index // 7) * 2
                date_cell = table.cell(date_row, col)
                self.add_cell_value(date_cell, str(d))
                shading_elm = parse_xml(r'<w:shd {} w:fill="D9D9D9"/>'.format(nsdecls('w')))
                date_cell._tc.get_or_add_tcPr().append(shading_elm)

                if not risk_count:
                    continue

                risk_row = 3 + (d_index // 7) * 2
                risk_cell = table.cell(risk_row, col)
                self.add_cell_value(risk_cell, str(risk_count))
                color = self.get_cell_background_color(risk_level)
                if color is not None:
                    risk_cell._tc.get_or_add_tcPr().append(color)
        for row in table.rows:
            row.height_rule = WD_ROW_HEIGHT_RULE.EXACTLY
            row.height = Mm(5)
        return sub_doc

    def calendar_table(self, data, date_list):
        if not date_list:
            return {}
        date_list = [datetime.datetime.strptime(i, "%Y-%m-%d") for i in date_list]
        month_list = list(sorted({(d.year, d.month) for d in date_list}, key=lambda x: (x[0], x[1])))
        if len(month_list) > 2:
            return self.year_calendar(month_list, data)
        else:
            return self.month_calendar(month_list, data)
